import torch
import torchvision
from torch.utils import data
from torchvision import transforms
from d2l import torch as d2l
d2l.use_svg_display()

'''
读取数据集，用的是torchvision的datasets包,里面包含了，可以直接下载至内存
# ToTensor可以将图片PIL类型transforms为Tensor，注意dataset里面的参数transform没有s！！！！
'''
mnist_train = torchvision.datasets.FashionMNIST(root='dl-learn-pytorch/data', train=True, transform=transforms.ToTensor(), download=True)
mnist_test = torchvision.datasets.FashionMNIST(root='dl-learn-pytorch/data', train=False, transform=transforms.ToTensor(), download=True)
len(mnist_train)
len(mnist_test)
mnist_train.shape # 测试数据和单个图片的大小
mnist_train[0][0].shape
# Fasionmnist包含十个类别，以下函数是数字索引到文字之间的转换
def get_fashion_mnist_label(labels):
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker',
                    'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels] #???，是否可以直接用return text_labels[int(i)]呢，答案是不行

# 以下是个人因为不熟悉python的for进行的测试
test_for1 = ['a', 'b', 'c', 'd']
test_for1[0]
test_labels = torch.tensor([0, 0, 2])
for i in test_labels:
    print(test_for1[int(i)]) # 输出结果为：a a c
def test_for(x):
    return [test_for1[int(x)] for x in test_labels]
def test_for_direct(x):
    return test_for1[int(x)] # 报错：不接受一个列表的返回，所以需要用for进行提取
test_for(test_for1) # 结果也是：a a c
test_for_direct(test_for1)

# 可视化这些样本
def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):
#@save """绘制图像列表"""
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            # 图片张量
            ax.imshow(img.numpy())
    else:
        # PIL图片
            ax.imshow(img)
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)
    if titles:
        ax.set_title(titles[i])
    return axes
X, y = next(iter(data.DataLoader(mnist_train, batch_size=18)))
show_images(X.reshape(18, 28, 28), 2, 9, titles=get_fashion_mnist_label(y));

# 读取小批量数据集
batch_size = 256
def get_data_workers(): # 使用4个进程来读取数据
    return 4
# 通过内置数据迭代器DataLoader，我们可 以随机打乱了所有样本，从而无偏见地读取小批量
train_iter = data.DataLoader(mnist_train, batch_size, shuffle=True, num_workers=get_data_workers())
# 看一下读取数据训练的时间
timer = d2l.Timer()
for X, y in train_iter:
    continue
f'{timer.stop():.2f} sec'

# 整合以上组件，构造一个下载和读取训练、测试数据集的loader迭代器。此外，这个函数还接受一个可选参数resize，用来将图像大小调整为另一种形状，
def load_data_fashion_mnist(batch_size, resize=None):
    trans = [transforms.ToTensor()] # 实例化一个张量的trans
    if resize: # 函数默认resize为None，如果不为空则执行调整图像操作
        trans.insert(0, transforms.Resize(resize))
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(root='dl-learn-pytorch/data_resize', train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(root='dl-learn-pytorch/data_resize', train=False, transform=trans, download=True)
    return (data.DataLoader(mnist_train, batch_size, shuffle=True, num_workers=get_data_workers()),
            data.DataLoader(mnist_test, batch_size, shuffle=False, num_workers=get_data_workers()))
# 测试整合组件
train_iter, test_iter = load_data_fashion_mnist(32, resize=64)
for X, y in train_iter:
    print(X.shape, X.dtype, y.shape, y.dtype)
    break

ac = torch.randn(2, 1, 4, 3) # 个人测试，高维张量形状判断，从后往前看，依次乘扩大即可